#include "AMGPCGsolver.h"
#include <deal.II/lac/sparsity_pattern.h>

void Gauss_Seidel(const DealiiSpMat &M, DealiiVectorXd &x,
                  const DealiiVectorXd &rhs, const int s) {
    // const SparsityPattern &spM = M.get_sparsity_pattern();
    // const std::size_t *rowstart = spM.get_rowstart_indices();
    // const u_int *colnums = spM.get_column_numbers();
    // for (u_int i = 0; i < s; i++) {
    //     for (u_int j = 0; j < M.m(); j++) {
    //         double r0 = rhs(j);
    //         for (u_int k = rowstart[j] + 1; k < rowstart[j + 1]; k++) {
    //             r0 -= x(colnums[k]) * M.global_entry(k);
    //         }
    //         x(j) = r0 / M.global_entry(rowstart[j]);
    //     }
    // }
    for(u_int i = 0; i < s; i++) {
      for (u_int j = 0; j < M.m(); j++) {
        double r0 = rhs(j);
        auto end = M.end(j);
        for(auto it = M.begin(j)+1; it != end;it++) {
          u_int col = it->column();
          double val = it->value();
          r0 -= x(col)*val;
        }
        double demon = M.begin(j)->value();
        x(j) = r0 / demon;
      }
    }
};

double KerError(const DealiiSpMat &spt) {
    DealiiVectorXd vec_one(spt.m());
    for (int i = 0; i < vec_one.size(); i++)
        vec_one[i] = 1;
    // std::cout <<"Vec_ONe"<< vec_one.l2_norm() << std::endl;
    DealiiVectorXd Au(spt.m());
    spt.vmult(Au, vec_one);
    return Au.l2_norm();
}

double L2_Error(DealiiVectorXd &v1, DealiiVectorXd &v2) {
    double error1 = 0;
    double error2 = 0;
    for (int i = 0; i < v1.size(); i++) {
        error1 += v1[i] * v1[i] + v2[i] * v2[i];
        error2 += v1[i] * v2[i];
    }
    return sqrt(error1 - 2 * error2) / v1.size();
};
void BuildMesh::build() {
    const SparsityPattern &spM = (*M).get_sparsity_pattern();
    // const std::size_t *M_rowstart = spM.get_rowstart_indices();
    // const u_int *M_colnums = spM.get_column_numbers();

    u32 n_core_point = 0;
    u32 n = (*M).m();
    vector<int> flag(n, 0);
    vector<double> beta(n, 0);
    for (u32 i = 0; i < n; i++) {
        beta[i] = 0.0;
        // for (u_int j = M_rowstart[i] + 1; j < M_rowstart[i + 1]; j++)
        //     beta[i] = std::max(fabs((*M).global_entry(j)), beta[i]);
        for(auto it = (*M).begin(i)+1; it != (*M).end(i); it++) {
          beta[i] = std::max(fabs(it->value()),beta[i]);
        }
        beta[i] *= theta;
        if (flag[i] != 0)
            continue;
        // for (u_int j = M_rowstart[i] + 1; j < M_rowstart[i + 1]; j++) {
        //     const u_int &k = M_colnums[j];
        //     if (flag[k] == -1)
        //         continue;
        //     if (fabs((*M).global_entry(j)) < beta[i])
        //         continue;
        //     flag[k] = 1;
        // }
        for(auto it = (*M).begin(i)+1; it != (*M).end(); it++) {
          const u_int& k = it->column();
          if(flag[k] == -1)
            continue;
          if(fabs(it->value()) < beta[i])
            continue;
          flag[k] = 1;
        }
        flag[i] = -1;
        n_core_point++;
    }
    // std::cout << n_core_point << std::endl;
    for (u_int i = 0; i < n; i++) {
        if (flag[i] != 1)
            continue;
        list<u_int> dep_pnt;
        // for (u_int j = M_rowstart[i] + 1; j < M_rowstart[i + 1]; j++) {
        //     const u_int &k = M_colnums[j];
        //     if (fabs((*M).global_entry(j)) < beta[i])
        //         continue;
        //     if (flag[k] == -1)
        //         dep_pnt.push_back(k);
        // }
        for(auto it = (*M).begin(i)+1; it != (*M).end(i); it++) {
          const u_int& k = it->column();
          if(fabs(it->value()) < beta[i])
            continue;
          if(flag[k] == -1)
            dep_pnt.push_back(k);
        }
        bool is_FF = false;
        // for (u_int j = M_rowstart[i] + 1; j < M_rowstart[i + 1]; j++) {
        //     const u_int &k = M_colnums[j];
        //     if ((flag[k] == -1) || (k < i) ||
        //         (fabs((*M).global_entry(j)) < beta[i]))
        //         continue;
        //     is_FF = true;
        //     for (u_int j1 = M_rowstart[k] + 1; j1 < M_rowstart[k + 1]; j1++) {
        //         const u_int &k1 = M_colnums[j1];
        //         if (flag[k1] != -1)
        //             continue;
        //         list<u32>::iterator result =
        //             std::find(dep_pnt.begin(), dep_pnt.end(), k1);
        //         if (result != dep_pnt.end()) {
        //             is_FF = false;
        //             break;
        //         }
        //     }
        //     if (is_FF)
        //         break;
        // }
        for(auto it = (*M).begin(i)+1; it != (*M).end(i); it++) {
          const u_int& k = it->column();
          if ((flag[k] == -1) || (k < i) ||
            (fabs(it->value()) < beta[i]))
              continue;
          is_FF = true;   
          for(auto it2 = (*M).begin(k)+1; it2 != (*M).end(k); it2++)
          {
            const u_int &k1 = it2->column();
            if (flag[k1] != -1)
              continue;
            list<u32>::iterator result =
              std::find(dep_pnt.begin(), dep_pnt.end(), k1);
            if (result != dep_pnt.end()) {
              is_FF = false;
              break;
            }
          }
            if (is_FF)
                break;
        } 
        if (!is_FF)
            continue;
        flag[i] = -1;
        n_core_point++;
    }

    Cpoints.clear();
    Cpoints.reserve(n_core_point);
    for (u32 i = 0; i < n; i++)
        if (flag[i] == -1)
            Cpoints.push_back(i);
    std::vector<u_int> core_point(n_core_point);
    std::vector<int> core_at((*M).m(), -1);

    vector<int> f2c(n, -1); //细网格点到粗网格点的索引映射
    for (u32 i = 0; i < Cpoints.size(); i++)
        f2c[Cpoints[i]] = i;

    collength.resize((*M).m());
    rowlength.resize(n_core_point);
    for (u32 i = 0; i < n; i++) {
        if (f2c[i] != -1) {
            collength[i]++;
            rowlength[f2c[i]]++;
            continue;
        }
        // for (u_int j = M_rowstart[i] + 1; j < M_rowstart[i + 1]; j++) {
        //     const u_int &k = M_colnums[j];
        //     if (f2c[k] != -1) {
        //         collength[i]++;
        //         rowlength[f2c[k]]++;
        //     }
        // }
        for (auto it = (*M).begin(i)+1; it != (*M).end(i); it++) {
          const u_int& k = it->column();
          if(f2c[k] != -1) {
            collength[i]++;
            rowlength[f2c[k]]++;
          }
        }
    }

    // for (u_int i = 0, n = 0;i < (*M).m();i ++) {
    //     if (flag[i] != -1) continue;

    //     core_point[n] = i;
    //     core_at[i] = n; //!
    //     for (u_int j = M_rowstart[i];j < M_rowstart[i+1];j ++) {
    //     const u_int& k = M_colnums[j];
    //     rowlength[n] ++;
    //     collength[k] ++;

    //     }
    //     n ++;
    // }
    // std::cout << n_core_point << std::endl;
    ///////////////粗网格点已经选择完毕，接下来进行插值算子的构建/////////////////

    SpItp.reinit(n, Cpoints.size(), collength);
    for (u32 i = 0; i < n; i++) {
        if (f2c[i] != -1) {
            SpItp.add(i, f2c[i]);
            continue;
        }
        // for (u_int j = M_rowstart[i] + 1; j < M_rowstart[i + 1]; j++) {
        //     const u_int &k = M_colnums[j];
        //     if (f2c[k] != -1)
        //         SpItp.add(i, f2c[k]);
        // }
        for (auto it = (*M).begin(i)+1; it != (*M).end(i); it++) {
          const u_int& k = it->column();
          if(f2c[k] != -1) {
            SpItp.add(i,f2c[k]);
          }
        }
    }
    SpItp.compress();
    ItplatMat.reinit(SpItp);
    for (u32 i = 0; i < n; i++) {
        if (f2c[i] != -1) {
            ItplatMat.add(i, f2c[i], 1.0);
            continue;
        }
        double SumAin = (*M).diag_element(i);
        vector<u32> Fs;
        vector<u32> CI;
        // for (u_int j = M_rowstart[i] + 1; j < M_rowstart[i + 1]; j++) {
        //     const u_int &k = M_colnums[j];
        //     if (f2c[k] != -1)
        //         CI.push_back(k);
        //     else {
        //         if (fabs((*M).global_entry(j)) < beta[i])
        //             SumAin += (*M).global_entry(j);
        //         else
        //             Fs.push_back(k);
        //     }
        // }
        for (auto it = M->begin(i)+1; it != M->end(i); it++) {
          const u_int &k = it->column();
          if (f2c[k] != -1)
            CI.push_back(k);
          else {
            if(fabs(it->value()) < beta[i])
              SumAin += it->value();
            else
              Fs.push_back(k);
          }
        }
        vector<double> SumAmk(Fs.size(), 0);
        for (u32 j = 0; j < Fs.size(); ++j) {
            u32 m = Fs[j];
            // for (u_int jj = M_rowstart[m]; jj < M_rowstart[m + 1]; jj++) {
            //     const u_int &row = M_colnums[jj];
            //     if (find(CI.begin(), CI.end(), row) != CI.end())
            //         SumAmk[j] += (*M).global_entry(jj);
            // }
            for (auto it = M->begin(m); it != M->end(m); it++) {
              const u_int &row = it->column();
              if (find(CI.begin(),CI.end(),row) != CI.end())
                SumAmk[j] += it->value();
            }
        }
        for (u32 j = 0; j < Fs.size(); ++j) {
            u32 m = Fs[j];
            // for (u_int jj = M_rowstart[m]; jj < M_rowstart[m + 1]; jj++) {
            //     const u_int &row = M_colnums[jj];
            //     if (find(CI.begin(), CI.end(), row) != CI.end())
            //         if (std::fabs(SumAmk[j]) > 1e-12)
            //             ItplatMat.add(i, f2c[row],
            //                           -1 * ((*M).el(i, m)) *
            //                               ((*M).global_entry(jj)) /
            //                               (SumAmk[j] * SumAin));
            // }
            for (auto it = M->begin(m)+1; it != M->end(m); it++) {
                const u_int &row = it->column();
                if (find(CI.begin(), CI.end(), row) != CI.end())
                    if (std::fabs(SumAmk[j]) > 1e-12)
                        ItplatMat.add(i, f2c[row],
                                      -1 * ((*M).el(i, m)) *
                                          (it->value()) /
                                          (SumAmk[j] * SumAin));
            }
        }
        for (auto &v : CI)
            ItplatMat.add(i, f2c[v], -1 * ((*M).el(i, v)) / SumAin);
    }
}

AMG::AMG(DealiiSpMat &_Spt, double _Tol, double _theta, int _smooth,
         int _maxlayer) {
    cout << "Begin to build mesh " << endl;
    Tolerance = _Tol;
    Theta = _theta;
    smooth_num = _smooth;
    MaxLayerNum = _maxlayer;
    N_layer = 1;
    double EndRatio = 0.25;
    ////注意这里重新分配空间会导致引用失效, 所以初始化就得分配好空间
    Vec_Mat.reserve(MaxLayerNum);
    Interpolation.reserve(MaxLayerNum);
    Restrict.reserve(MaxLayerNum);
    Vec_SPM.reserve(MaxLayerNum);
    Vec_SPP.reserve(MaxLayerNum);
    Vec_SPPt.reserve(MaxLayerNum);
    auto t_start = std::chrono::high_resolution_clock::now();
    ///////////矩阵的复制//////////
    {
        DealiiSpMat mm;
        SparsityPattern spmm;
        Vec_Mat.push_back(mm);
        // const SparsityPattern &spM = _Spt.get_sparsity_pattern();
        // const std::size_t *rowstart = spM.get_rowstart_indices();
        // const u_int *colnums = spM.get_column_numbers();
        //
        // Vec_SPM.push_back(spM);
        // Vec_SPM.back().copy_from(spM);
        // Vec_Mat.back().reinit(Vec_SPM.back());
        // for (u_int j = 0; j < _Spt.m(); j++) {
        //     for (u_int k = rowstart[j]; k < rowstart[j + 1]; k++) {
        //         Vec_Mat.back().add(j, colnums[k], _Spt.global_entry(k));
        //     }
        // }
        const SparsityPattern& spM = _Spt.get_sparsity_pattern();
        Vec_SPM.push_back(spmm);
        Vec_SPM.back().copy_from(spM);
        Vec_Mat.back().reinit(Vec_SPM.back());
        // Vec_Mat.back().copy_from(_Spt);
        for(u32 i = 0; i < _Spt.m(); i++) {
          for(auto it = _Spt.begin(i); it != _Spt.end(i); it++) {
            Vec_Mat.back().add(i,it->column(),it->value());
          }
        }
    }
    while (1.0 * Vec_Mat.back().n_nonzero_elements() <
               EndRatio * Vec_Mat.back().m() * Vec_Mat.back().m() &&
           N_layer <= MaxLayerNum && Vec_Mat.back().n() > 500) {
        // cout << Vec_Mat.back().m() << " ";
        // cout << 1.0 * Vec_Mat.back(). n_nonzero_elements() /
        // (Vec_Mat.back().m() * Vec_Mat.back().m()) << " " <<
        // KerError(Vec_Mat.back()) <<  endl;
        N_layer++;
        BuildMesh bdmesh(Vec_Mat.back(), Theta);
        bdmesh.build();
        DealiiSpMat &itp = bdmesh.GetMatP();
        /////////插值算子和投影算子
        vector<u32> &length = bdmesh.Getlength();
        SparsityPattern Dsp(itp.n(), itp.m(), length);
        const SparsityPattern &sp = itp.get_sparsity_pattern();
        // const std::size_t *rowstart = sp.get_rowstart_indices();
        // const u_int *colnums = sp.get_column_numbers();
        Interpolation.push_back(itp);
        Restrict.push_back(_Spt);
        Vec_SPPt.push_back(sp);
        Vec_SPPt.back().copy_from(sp);
        Interpolation.back().reinit(Vec_SPPt.back());
        for (u_int j = 0; j < itp.m(); j++) {
            // for (u_int k = rowstart[j]; k < rowstart[j + 1]; k++) {
            //     Dsp.add(colnums[k], j);
            //     Interpolation.back().add(j, colnums[k], itp.global_entry(k));
            // }
            for (auto it = itp.begin(j); it < itp.end(j); it++) {
              Dsp.add(it->column(),j);
              Interpolation.back().add(j,it->column(),it->value());
            }
        }
        Dsp.compress();
        Vec_SPP.push_back(Dsp);
        Vec_SPP.back().copy_from(Dsp);
        Restrict.back().reinit(Vec_SPP.back());
        for (u_int j = 0; j < itp.m(); j++) {
            // for (u_int k = rowstart[j]; k < rowstart[j + 1]; k++) {
            //     Restrict.back().add(colnums[k], j, itp.global_entry(k));
            // }
            for (auto it = itp.begin(j); it != itp.end(); it++) {
              Restrict.back().add(it->column(),j,it->value());
            }
        }
        ////////////三连乘PMPT///////////////////
        DealiiSpMat &P = Restrict.back();
        DealiiSpMat &M = Vec_Mat.back();
        DealiiSpMat &Pt = Interpolation.back();
        ////连续矩阵乘法
        SparsityPattern PMsp;
        DealiiSpMat PM;
        PM.reinit(PMsp);
        P.mmult(PM, M);
        SparsityPattern Asp;
        Vec_Mat.push_back(_Spt);
        Vec_SPM.push_back(PMsp);
        Vec_Mat.back().reinit(Vec_SPM.back());
        PM.mmult(Vec_Mat.back(), Pt);
    }
    // std::cout << Vec_Mat.back().m() <<" \n" ;
    //////////////////最粗网格使用直接法////////////////
    A_direct.initialize(Vec_Mat.back());
    F_Rhs.resize(N_layer);
    U_Lhs.resize(N_layer);
    auto t_end = std::chrono::high_resolution_clock::now();
    double t_total =
        std::chrono::duration<float, std::milli>(t_end - t_start).count();
    cout << "Build Matrices over! It takes  " << t_total << " .ms" << endl;
}

void AMG::Vcycle(int _layer) {
    if (_layer == N_layer - 1) {
        U_Lhs[_layer] = F_Rhs[_layer];
        A_direct.solve(U_Lhs[_layer]);
        return;
    } else {
        DealiiVectorXd &u = U_Lhs[_layer];
        DealiiSpMat &A = Vec_Mat[_layer];
        DealiiVectorXd &f = F_Rhs[_layer];
        Gauss_Seidel(A, u, f, smooth_num);
        DealiiVectorXd r_h = f;
        DealiiVectorXd Au(f.size());
        A.vmult(Au, u);
        r_h -= Au;
        F_Rhs[_layer + 1].reinit(Restrict[_layer].m());
        Restrict[_layer].vmult(F_Rhs[_layer + 1], r_h);
        U_Lhs[_layer + 1].reinit(F_Rhs[_layer + 1].size());
        Vcycle(_layer + 1);
        DealiiVectorXd e_h(u.size());
        Interpolation[_layer].vmult(e_h, U_Lhs[_layer + 1]);
        u += e_h;
        Gauss_Seidel(A, u, f, smooth_num);
    }
};

void AMG::Vcycle() { Vcycle(0); }
DealiiVectorXd &AMG::Vcycle(const DealiiVectorXd &Rhs) {
    U_Lhs.resize(N_layer);
    F_Rhs.resize(N_layer);
    U_Lhs[0].reinit(Rhs.size());
    F_Rhs[0] = Rhs;
    Vcycle();
    return U_Lhs[0];
}
DealiiVectorXd &AMG::solve(const DealiiVectorXd &Rhs) {
    int N_vcycle = 0;
    U_Lhs.resize(N_layer);
    F_Rhs.resize(N_layer);
    U_Lhs[0].reinit(Rhs.size());
    F_Rhs[0] = Rhs;
    DealiiSpMat &A = Vec_Mat[0];
    DealiiVectorXd &u = U_Lhs[0];
    DealiiVectorXd e_init(Rhs);
    DealiiVectorXd rhs = Rhs;
    DealiiVectorXd Au(Rhs.size());
    A.vmult(Au, u);
    e_init -= Au;
    double initerror = Rhs.l2_norm();
    double error = initerror;
    cout << "Init error is " << initerror << endl;
    while (error > initerror * Tolerance && N_vcycle < 200) {
        Vcycle(0);
        e_init = Rhs;
        A.vmult(Au, U_Lhs[0]);
        e_init -= Au;
        N_vcycle++;
        error = e_init.l2_norm();
        // cout << error << endl;
    }
    // double initerror = L2_Error(Au, rhs);
    // double Tol = Tolerance * initerror;
    // cout << "Init error is "<< initerror << endl;
    // double error = initerror;
    // while (error > Tol && N_vcycle < 200)
    // {
    //     Vcycle(0);
    //     A.vmult(Au, U_Lhs[0]);
    //     error = L2_Error(Au, rhs);
    //     N_vcycle++;
    //     // cout << error << endl;
    // }
    cout << N_vcycle << " th iterator, error is " << error << endl;
    return U_Lhs[0];
}
void RebuildMat(DealiiSpMat &M) {
    const SparsityPattern &spM = M.get_sparsity_pattern();
    // const std::size_t *rowstart = spM.get_rowstart_indices();
    // const u_int *colnums = spM.get_column_numbers();
    for (u_int j = 0; j < M.m(); j++) {
        double val = M.diag_element(j); //每一行非对角元的正元和
        // for (u_int k = rowstart[j] + 1; k < rowstart[j + 1]; k++) {
        //     if (M.global_entry(k) > 0) {
        //         val += M.global_entry(k);
        //         M.set(j, colnums[k], 0);
        //     }
        //     M.set(j, j, val);
        // }
        for (auto it = M.begin(j)+1; it != M.end(j); it++) {
          if(it->value() > 0) {
            val += it->value();
            M.set(j,it->column(),0);
          }
          M.set(j,j,val);
        }
    }
}
void AMG::Vcycleonce(DealiiVectorXd &x, const DealiiVectorXd &Rhs) {
    U_Lhs[0] = x;
    F_Rhs[0] = Rhs;
    Vcycle();
    x = U_Lhs[0];
}
void AMGPCG::BuildAMG() {
    DealiiSpMat M;
    SparsityPattern spM; //(SpA)
    spM.copy_from(SpA);
    // const std::size_t *rowstart = SpA.get_rowstart_indices();
    // const u_int *colnums = SpA.get_column_numbers();
    // M.reinit(spM);
    // for (u_int j = 0; j < A.m(); j++) {
    //     for (u_int k = rowstart[j]; k < rowstart[j + 1]; k++) {
    //         M.add(j, colnums[k], A.global_entry(k));
    //     }
    // }
    M.reinit(spM);
    //M.copy_from(A);
    for(u32 i = 0; i < A.m(); i++) {
      for(auto it = A.begin(i); it != A.end(i); it++){
        M.add(i,it->column(),it->value());
      }
    }
    if (rebuild)
        RebuildMat(M);
    AmgGrid = new AMG(M, Tolerance, theta, Smoothstep);
}
DealiiVectorXd AMGPCG::solve(DealiiVectorXd &_Rhs) {
    std::cout << "begin to solve equ by AMGPcg!" << std::endl;
    if (_Rhs.l2_norm() < Tolerance)
        return _Rhs;
    Rhs = _Rhs;
    int k = 0;
    xk.reinit(Rhs.size());
    DealiiVectorXd veczero(Rhs.size());

    double AmgTimeCount = 0.0;
    double CgTimeCount = 0.0;
    /// @brief 先做几次迭代，以获得初始猜测值x0
    /// @return
    Gauss_Seidel(A, xk, Rhs, 2);

    rk = Rhs;
    DealiiVectorXd Au(rk.size());
    A.vmult(Au, xk);
    rk -= Au;

    // double initerror = L2_Error(r0, veczero);
    double initerror = rk.l2_norm();
    cout << "Begin cg iterator \n"
         << "init error is " << initerror << endl;
    // while (L2_Error(rk[k], veczero) > Tolerance * initerror)
    while (rk.l2_norm() > Tolerance * initerror && k < Max_itenum) {
        auto amgt_start = std::chrono::high_resolution_clock::now();
        zk = AmgGrid->Vcycleonce(rk);

        auto amgt_end = std::chrono::high_resolution_clock::now();
        double amgtcount =
            std::chrono::duration<double, std::milli>(amgt_end - amgt_start)
                .count();
        AmgTimeCount += amgtcount;
        ++k;
        if (k == 1)
            pk = zk;
        else {
            betak = rk * zk / (rk_1 * zk_1);
            pk = zk;
            pk.add(betak, pk_1);
        }
        rk_1 = rk;
        zk_1 = zk;
        pk_1 = pk;
        DealiiVectorXd Apk(xk.size());
        A.vmult(Apk, pk);
        alphak = rk * zk / (pk * Apk);

        xk.add(alphak, pk);

        rk.add(-1 * alphak, Apk);
        // cout << k << "th iterator , error is " << rk.l2_norm() << endl;
    }
    // cout << k << "th iterator , error is " << L2_Error(rk[k], veczero) <<
    // endl;
    if (k >= Max_itenum) {
        cout << k << "th iterator , error is " << rk.l2_norm() << endl;
        cout << "Can't convergence! " << std::endl;
        return xk;
    }

    cout << k << "th iterator , AMG takes " << AmgTimeCount << ".ms" << endl;
    A.vmult(Au, xk);
    rk = Rhs;
    rk -= Au;
    // cout << "L2 error is " << L2_Error(r0, veczero) << endl;
    cout << "L2 error is " << rk.l2_norm() << endl;
    return xk;
}
